from src.geometry import LineSegment, Collision
from .node import Node
from .rrt_map import RRTMap
from .rrt_planner import RRTPlanner
from .robot_rrt_parameter import RobotRRTParameter
from src.geometry import Distance

from typing import List
from ..path_parameter import PathParameter
from ..path_planner import PathPlanner
from ..joint_planning import JointParameter, JointPlanner

import numpy as np


class RobotRRTPlanner(RRTPlanner):
    def __init__(self, rrt_map: RRTMap, rrt_parameter: RobotRRTParameter) -> None:
        self.robot = rrt_parameter.robot
        super().__init__(rrt_map, rrt_parameter)

    def check_collision(self, line_segment: LineSegment) -> bool:
        q0 = line_segment.get_point0().get_t()
        q1 = line_segment.get_point1().get_t()

        count = 10
        for i in range(count + 1):
            q = q0 + (q1 - q0) / count * i
            self.robot.set_joint(q)

            geometries = self.robot.get_geometries()
            for obstacles in self.obstacles:
                for geometry in geometries:
                    if Collision.is_collision(obstacles, geometry):
                        return True
        return False

    def get_path_parameters(self) -> List[PathParameter]:
        joint_parameters = []
        points = self.path[::-1]
        for i, point in enumerate(points[:-1]):
            joint_parameters.append(JointParameter(point.get_t(), points[i + 1].get_t()))
        return joint_parameters

    def interpolate(self, s: float):
        path_parameters = self.get_path_parameters()
        total_distance = s * self.last_path_length
        index = round(total_distance // self.expand_dis)
        local_distance = total_distance % self.expand_dis / np.linalg.norm(
            path_parameters[index].get_t1() - path_parameters[index].get_t0()
        )

        path_planner = PathPlanner(path_parameters[index])
        return path_planner.interpolate(local_distance)
